#include "imu_factor.h"

void ImuFactor::ComputeResidualAndJacobianPoseGraph(const State& last_imu_state, const State& imu_state, const State& last_state, const State& state, Eigen::Matrix<double, 5, 1> &residual, Eigen::Matrix<double, 5, 8> &jacobianXi, Eigen::Matrix<double, 5, 5> &jacobianXj, const Eigen::Vector3d& last_imu, const Eigen::Vector3d& curr_imu, const double& dt) {
	Eigen::Matrix3d imu_obs_pose_i = math_utils::compute_pose(Eigen::Vector3d(last_imu_state.p[0],last_imu_state.p[1],last_imu_state.theta));
	Eigen::Matrix3d imu_obs_pose_j = math_utils::compute_pose(Eigen::Vector3d(imu_state.p[0],imu_state.p[1],imu_state.theta));
	Eigen::Matrix3d imu_obs_pose_ij = imu_obs_pose_i.inverse() * imu_obs_pose_j;
	Eigen::Matrix2d Rij_obs = imu_obs_pose_ij.block<2,2>(0,0);
	Eigen::Matrix2d Rij_t_obs = Rij_obs.inverse();
	Eigen::Vector2d tij_obs = imu_obs_pose_ij.block<2,1>(0,2);
	Eigen::Matrix2d Ri = math_utils::rotation(last_state.theta);
	Eigen::Vector2d pi = last_state.p;
	Eigen::Matrix2d Rj = math_utils::rotation(state.theta);
	Eigen::Vector2d pj = state.p;

	residual.block<2,1>(0,0) = Rij_t_obs * (Ri.transpose()*(pj-pi)-tij_obs);
	residual.block<2,1>(2,0) = state.v - last_state.v - (imu_state.v - last_imu_state.v);
	residual(4,0) = state.theta - last_state.theta - (imu_state.theta - last_imu_state.theta);

	// cout << "ImuFactor residual =========== : " << residual.transpose() << endl;
	
	// compute jacobian
	jacobianXi.block<2,2>(0,0) = -Rij_t_obs * Ri.transpose(); // (1)
	jacobianXi.block<2,2>(0,2) = state.dt * Eigen::Matrix2d::Identity(); // (2)

	Eigen::Matrix<double,2,2> sign = Eigen::Matrix2d::Zero();
    sign(0,1) = -1;
    sign(1,0) = 1; 
    Eigen::Matrix<double,2,2> Ri_t_dev = sign * Ri.transpose();

	jacobianXi.block<2,1>(0,4) = Rij_t_obs * Ri_t_dev * (pj-pi); // (3)
	jacobianXi.block<2,2>(0,5) = 0.5 * Ri * state.dt * state.dt; // (4)
	jacobianXi.block<2,1>(0,7) = Eigen::Matrix<double,2,1>::Zero(); // (5)

	jacobianXi.block<2,2>(2,0) = Eigen::Matrix2d::Zero(); // (17)
	jacobianXi.block<2,2>(2,2) = -Eigen::Matrix2d::Identity(); // (18)

	Eigen::Vector3d last_imu_unbias(last_imu[0], last_imu[1], last_imu[2]);
    Eigen::Vector3d curr_imu_unbias(curr_imu[0], curr_imu[1], curr_imu[2]);

    last_imu_unbias[0] -= last_state.ba[0];
    curr_imu_unbias[0] -= last_state.ba[0];
    last_imu_unbias[1] -= last_state.ba[1];
    curr_imu_unbias[1] -= last_state.ba[1];
    last_imu_unbias[2] -= last_state.bg;
    curr_imu_unbias[2] -= last_state.bg;
    
    // std::cout << "dt ============= : " << dt << ", " << init_bias_ << std::endl;
    Eigen::Vector3d imu_unbias_hat = 0.5 * (last_imu_unbias + curr_imu_unbias);
    Eigen::Matrix<double,2,2> Ri_dev = sign * Ri;
    
	jacobianXi.block<2,1>(2,4) = dt * Ri_dev * imu_unbias_hat.block<2,1>(0,0); // (19)
	jacobianXi.block<2,2>(2,5) = dt * Ri; // (20)
	jacobianXi.block<2,1>(2,7) = Eigen::Matrix<double,2,1>::Zero(); // (21)

	jacobianXi.block<1,2>(4,0) = Eigen::Matrix<double,1,2>::Zero(); // (6)
	jacobianXi.block<1,2>(4,2) = Eigen::Matrix<double,1,2>::Zero(); // (7)
	jacobianXi(4,4) = -1; // (8)
	jacobianXi.block<1,2>(4,5) = Eigen::Matrix<double,1,2>::Zero(); // (9)
	jacobianXi(4,7) = 1; // (10)

	jacobianXj.block<2,2>(0,0) = Rij_t_obs*Ri.transpose(); // (11)
	jacobianXj.block<2,2>(0,2) = Eigen::Matrix2d::Zero(); // (12)
	jacobianXj.block<2,1>(0,4) = Eigen::Matrix<double,2,1>::Zero(); // (13)

	jacobianXj.block<2,2>(2,0) = Eigen::Matrix<double,2,2>::Zero(); // (22)
	jacobianXj.block<2,2>(2,2) = Eigen::Matrix<double,2,2>::Identity(); // (23)
	jacobianXj.block<2,1>(2,4) = Eigen::Matrix<double,2,1>::Zero(); // (24)

	jacobianXj.block<2,2>(4,0) = Eigen::Matrix<double,2,2>::Zero(); // (14)
	jacobianXj.block<2,2>(4,2) = Eigen::Matrix<double,2,2>::Zero(); // (15)
	jacobianXj(4,4) = 1; // (16)
}